/**
  ******************************************************************************
  * Copyright (C) 2020 - ~, SCUT-RobotLab Development Team
  * @file    dynamicModel.cpp
  * @author  Mentos Seetoo
  * @brief   This file provides dynamic data structure containing parameters for 
  *          quadruped robot using floating base model of MIT Cheetah Code. 
  * 		 Kinematics/Invers Kinematics solutions are provided, too.			
  * @note    THe following third-party libraries are used:
  *          - STL
  *          - Eigen3
  * @warning 
  *     - At least C++11 is required.														 
  *     - This file refers to Quadruped.cpp' in MIT Cheetah Control code. 
  * 	- More info about frame definition plz view the document paper or webside:
  * 		http://www.scut-robotlab.cn
  ******************************************************************************
  * @attention
  * 
  * if you had modified this file, please make sure there is no any bug in your
  * code. Update the version Number and write dowm your name and the date.What 
  * most important is that make sure the users will have a clear understanding
  * through your description.
  ******************************************************************************
  */
//TODO detailed test is needed!
/* Includes ------------------------------------------------------------------*/
#include "Utilities/common_def.h"
#include "Dynamics/dynamicModel.h"
#include <math.h>
/* Founctions ----------------------------------------------------------------*/
using namespace std;
using namespace Quadruped;
using namespace CONST_DEF;

/**
 * @brief Build model by the dynamic parameters.
 * @return true build successfully
 * @return false failed to build
 */
template <typename T>
bool Dynamic<T>::buildModel()
{
	if (!paramIsValid())
		return false;

	const int baseID = 5;
	int bodyID = baseID;
	Mat3<T> I3 = Mat3<T>::Identity();
	/*
        We assumen robot's body (not including rotors) can be modeled as a
        uniformly distributed box.
    */
	Vec3<T> bodyDims(Param.BODY_LENGTH, Param.BODY_WIDTH, Param.BODY_HEIGHT);
	Model.addBase(_bodyInertia);
	Model.addGroundContactBoxPoints(5, bodyDims);

	/*
        Loop over 4legs
    */
	for (int legID(LF); legID < 4; legID++)
	{
		// Ab/Ad joint
		bodyID++;
		Mat6<T> xtreeAbad = createSXform(I3, convertWithLegSigns<T>(_abadLocation, legID));
		Mat6<T> xtreeAbadRotor = createSXform(I3, convertWithLegSigns<T>(_abadRotorLocation, legID));

		if (getLegSign(legID) == (int)RIGHT_SIDE)
		{
			Model.addBody(_abadInertia.flipAlongAxis(CoordinateAxis::Y),
						  _abadRotorInertia.flipAlongAxis(CoordinateAxis::Y),
						  Param.ABAD_GEAR_RATIO, baseID, JointType::Revolute,
						  CoordinateAxis::X, xtreeAbad, xtreeAbadRotor);
		}
		else
		{
			Model.addBody(_abadInertia, _abadRotorInertia, Param.ABAD_GEAR_RATIO, baseID,
						  JointType::Revolute, CoordinateAxis::X, xtreeAbad,
						  xtreeAbadRotor);
		}

		// Hip Joint
		bodyID++;
		Mat6<T> xtreeHip =
			createSXform(coordinateRotation<T>(CoordinateAxis::Z, T(M_PI)),
						 convertWithLegSigns<T>(_hipLocation, legID));
		Mat6<T> xtreeHipRotor =
			createSXform(coordinateRotation<T>(CoordinateAxis::Z, T(M_PI)),
						 convertWithLegSigns<T>(_hipRotorLocation, legID));
		if (getLegSign(legID) == (int)RIGHT_SIDE)
		{
			Model.addBody(_hipInertia.flipAlongAxis(CoordinateAxis::Y),
						  _hipRotorInertia.flipAlongAxis(CoordinateAxis::Y),
						  Param.HIP_GEAR_RATIO, bodyID - 1, JointType::Revolute,
						  CoordinateAxis::Y, xtreeHip, xtreeHipRotor);
		}
		else
		{
			Model.addBody(_hipInertia, _hipRotorInertia, Param.HIP_GEAR_RATIO, bodyID - 1,
						  JointType::Revolute, CoordinateAxis::Y, xtreeHip,
						  xtreeHipRotor);
		}

		// add knee ground contact point
		Model.addGroundContactPoint(bodyID, Vec3<T>(0, 0, -Param.HIP_LINK_LEN));

		// Knee Joint
		bodyID++;
		Mat6<T> xtreeKnee = createSXform(I3, _kneeLocation);
		Mat6<T> xtreeKneeRotor = createSXform(I3, _kneeRotorLocation);
		if (getLegSign(legID) == (int)RIGHT_SIDE)
		{
			Model.addBody(_kneeInertia.flipAlongAxis(CoordinateAxis::Y),
						  _kneeRotorInertia.flipAlongAxis(CoordinateAxis::Y),
						  Param.KNEE_GEAR_RATIO, bodyID - 1, JointType::Revolute,
						  CoordinateAxis::Y, xtreeKnee, xtreeKneeRotor);

			Model.addGroundContactPoint(bodyID, Vec3<T>(0, Param.KNEE_Y_OFFSET, -Param.KNEE_LINK_LEN), true);
		}
		else
		{
			Model.addBody(_kneeInertia, _kneeRotorInertia, Param.KNEE_GEAR_RATIO, bodyID - 1,
						  JointType::Revolute, CoordinateAxis::Y, xtreeKnee,
						  xtreeKneeRotor);

			Model.addGroundContactPoint(bodyID, Vec3<T>(0, -Param.KNEE_Y_OFFSET, -Param.KNEE_LINK_LEN), true);
		}
	}

	Vec3<T> g(0, 0, -9.81);
	Model.setGravity(g);

	return true;
}

/**
 * @brief  Check if the dynamics parameter is initialized. 
 * @return true 
 * @return false 
 */
template <typename T>
bool Dynamic<T>::paramIsValid()
{
	if (Param.BODY_LENGTH * Param.BODY_WIDTH * Param.BODY_HEIGHT *
			Param.BODY_MASS ==
		0)
		return false;
	if (Param.ABAD_LINK_LEN * Param.HIP_LINK_LEN * Param.KNEE_LINK_LEN == 0)
		return false;
	if (Param.ABAD_GEAR_RATIO * Param.HIP_GEAR_RATIO *
			Param.KNEE_GEAR_RATIO ==
		0)
		return false;
	if (Param.MOTOR_MAX_TORQUE * Param.MOTOR_KT * Param.MOTOR_R *
			Param.SUPPLY_VOLTAGE ==
		0)
		return false;

	return true;
}

/**
 * @brief Build actuator models for a leg
 * @tparam T 
 * @return std::vector<ActuatorModel<T>> 
 */
template <typename T>
bool Dynamic<T>::buildActuatorModels()
{
	if (!paramIsValid())
		return false;

	actuatorModels.emplace_back(Param.ABAD_GEAR_RATIO, Param.MOTOR_KT, Param.MOTOR_R, Param.SUPPLY_VOLTAGE,
								Param.JOINT_DAMP, Param.JOINT_FRICTION, Param.MOTOR_MAX_TORQUE);
	actuatorModels.emplace_back(Param.HIP_GEAR_RATIO, Param.MOTOR_KT, Param.MOTOR_R, Param.SUPPLY_VOLTAGE,
								Param.JOINT_DAMP, Param.JOINT_FRICTION, Param.MOTOR_MAX_TORQUE);
	actuatorModels.emplace_back(Param.KNEE_GEAR_RATIO, Param.MOTOR_KT, Param.MOTOR_R, Param.SUPPLY_VOLTAGE,
								Param.JOINT_DAMP, Param.JOINT_FRICTION, Param.MOTOR_MAX_TORQUE);
	return true;
}

/**
 * @brief Kinematics/Invers Kinematics. Transform between Body frame and Joint Space.
 * 		  We assume that angle of knee ranges from 0 - 180 degrees,
 * 		  angle of hip ranges from 0 - 180 degrees,
 * 	 	  angle of abdu/addu ranges from -90 to 90 degrees.
 * @param pFoot_B foot position relative to body frame
 * @param pLeg	  leg joints angle in joint space
 * @return true:success 	flase:bad input param
 */
template <typename T>
bool Dynamic<T>::convertToJointAngle(const Vec3<T> &pFoot_B, Vec3<T> &pLeg, int leg)
{
	T l1 = Param.ABAD_LINK_LEN;
	T l2 = Param.HIP_LINK_LEN;
	T l3 = Param.KNEE_LINK_LEN;

	//!< Cartesian coordinates relative to body frame.
	T x = pFoot_B(0), y = pFoot_B(1), z = pFoot_B(2);

	//!< Spherical coordinate relative to body frame.
	T alpha = getLegSign(leg) * atan2(sqrt(fabs(y * y + z * z - l1 * l1)) , l1) + atan2(y, -z) - M_PI_2;
	T r = sqrt(fabs(x * x + y * y + z * z - l1 * l1)); //!< rotation R between foot and hip.
	T beta = acos(LimitRange(x / r,(T)-1,(T)1)); //!< Assume r will never be zero

	//!< Inverse kinemamics resolved in geometry.
	pLeg(0) = alpha;
	pLeg(1) = acos(LimitRange((l2 * l2 + r * r - l3 * l3) / (2 * l2 * r), (T)-1, (T)1)) + beta;
	pLeg(2) = acos(LimitRange((l2 * l2 + l3 * l3 - r * r) / (2 * l2 * l3),(T)-1, (T)1));

	return true;
}

template <typename T>
bool Dynamic<T>::convertToJointAngle()
{
	T l1 = Param.ABAD_LINK_LEN;
	T l2 = Param.HIP_LINK_LEN;
	T l3 = Param.KNEE_LINK_LEN;

	for (uint8_t leg_x = 0; leg_x < LEG_NUM; leg_x++)
	{
		//!< Cartesian coordinates relative to body frame.
		T x = LegCommand[leg_x].p(0), y = LegCommand[leg_x].p(1), z = LegCommand[leg_x].p(2);

		//!< Spherical coordinate relative to body frame.
		T alpha = getLegSign(leg_x) * atan2(sqrt(fabs(y * y + z * z - l1 * l1)) , l1) + atan2(y, -z) - M_PI_2;
		T r = sqrt(fabs(x * x + y * y + z * z - l1 * l1)); //!< rotation R between foot and hip.
		T beta = acos(LimitRange(x / r,(T)-1,(T)1)); //!< Assume r will never be zero

		//!< Inverse kinemamics resolved in geometry.
		LegCommand[leg_x].q(0) = alpha;
		LegCommand[leg_x].q(1) = acos(LimitRange((l2 * l2 + r * r - l3 * l3) / (2 * l2 * r), (T)-1, (T)1)) + beta;
		LegCommand[leg_x].q(2) = acos(LimitRange((l2 * l2 + l3 * l3 - r * r) / (2 * l2 * l3),(T)-1, (T)1));
	}

	return true;
}

template <typename T>
bool Dynamic<T>::calFootPositions()
{
	T l1 = Param.ABAD_LINK_LEN;
	T l2 = Param.HIP_LINK_LEN;
	T l3 = Param.KNEE_LINK_LEN;

	T theta1, theta2, theta3, r, beta;
	for (uint8_t leg_x = 0; leg_x < LEG_NUM; leg_x++)
	{
		//!< Joint space discription.
		theta1 = LegState[leg_x].q(0);
		theta2 = LegState[leg_x].q(1);
		theta3 = LegState[leg_x].q(2);
		theta3 = theta3 < 0 ? 0 : theta3;

		//!< Spherical coordinate relative to body frame.
		r = sqrt(l3 * l3 + l2 * l2 - 2 * l3 * l2 * cos(theta3));
		beta = theta2 - acos((l2 * l2 + r * r - l3 * l3) / (2 * l2 * r));

		//!< Kinemamics solved in geometry.
		LegState[leg_x].p(0) = r * cos(beta);
		r *= sin(beta);			   //!< r projection in y-z plane
		r = sqrt(r * r + l1 * l1); //!< r of the foot in y-z plane
		theta1 += getLegSign(leg_x) * (M_PI_2 - acos(l1 / r));
		LegState[leg_x].p(1) = sin(theta1) * r;
		LegState[leg_x].p(2) = -cos(theta1) * r;

		//!< Calculate Jacobian
		T s1 = std::sin(LegState[leg_x].q(0));
		T s2 = std::sin(LegState[leg_x].q(1));
		T s3 = std::sin(LegState[leg_x].q(2));

		T c1 = std::cos(LegState[leg_x].q(0));
		T c2 = std::cos(LegState[leg_x].q(1));
		T c3 = std::cos(LegState[leg_x].q(2));

		T c23 = c2 * c3 - s2 * s3;
		T s23 = s2 * c3 + c2 * s3;

		LegState[leg_x].J_B.operator()(0, 0) = 0;
		LegState[leg_x].J_B.operator()(0, 1) = l3 * c23 + l2 * c2;
		LegState[leg_x].J_B.operator()(0, 2) = l3 * c23;
		LegState[leg_x].J_B.operator()(1, 0) = l3 * c1 * c23 + l2 * c1 * c2 - l1 * getLegSign(leg_x) * s1;
		LegState[leg_x].J_B.operator()(1, 1) = -l3 * s1 * s23 - l2 * s1 * s2;
		LegState[leg_x].J_B.operator()(1, 2) = -l3 * s1 * s23;
		LegState[leg_x].J_B.operator()(2, 0) = l3 * s1 * c23 + l2 * c2 * s1 + l1 * getLegSign(leg_x) * c1;
		LegState[leg_x].J_B.operator()(2, 1) = l3 * c1 * s23 + l2 * c1 * s2;
		LegState[leg_x].J_B.operator()(2, 2) = l3 * c1 * s23;
	}

	return true;
}

template <typename T>
bool Dynamic<T>::convertToCartesianSpace(const Vec3<T>& pLeg, Vec3<T> &pFoot_B)
{
	T l1 = Param.ABAD_LINK_LEN;
	T l2 = Param.HIP_LINK_LEN;
	T l3 = Param.KNEE_LINK_LEN;
	T theta1, theta2, theta3, r, beta;

	//!< Joint space discription.
	theta1 = pLeg(0);
	theta2 = pLeg(1);
	theta3 = pLeg(2);
	theta3 = theta3 < 0 ? 0 : theta3;

	//!< Spherical coordinate relative to body frame.
	r = sqrt(l3 * l3 + l2 * l2 - 2 * l3 * l2 * cos(theta3));
	beta = theta2 - acos((l2 * l2 + r * r - l3 * l3) / (2 * l2 * r));

	//!< Kinemamics solved in geometry.
	pFoot_B.operator()(0) = r * cos(beta);
	r *= sin(beta);			   //!< r projection in y-z plane
	r = sqrt(r * r + l1 * l1); //!< r of the foot in y-z plane
	theta1 += (theta1 > 0 ? 1 : -1) * (M_PI_2 - acos(l1 / r));
	pFoot_B.operator()(1) = sin(theta1) * r;
	pFoot_B.operator()(2) = -cos(theta1) * r;

	return true;
}

template class Dynamic<float>;
template class Dynamic<double>;

/**
 * @brief Get if the leg is on the left(+) or right(-) of the robot
 * @param leg : leg index 
 * @return 1 for left, -1 for right.
 */
int getLegSign(int legID)
{
	const int sign[2] = {1, -1};
	return (legID == LF || legID == LB) ? sign[0] : sign[1];
}
/************************ COPYRIGHT(C) SCUT-ROBOTLAB **************************/